
## Snakes on a Payload
## Copyright (C) 2011 CORE Security Technologies
##
## This program is free software; you can redistribute it and/or
## modify it under the terms of the GNU General Public License
## as published by the Free Software Foundation; either version 2
## of the License, or (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program; if not, write to the Free Software
## Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
##
## Authors:
## 2011/09/11 - Pedro Varangot (pvarangot@)
## 2011/09/11 - Fernando Russ (fruss@)



import subprocess
import re

from relocation import Relocation
from symbol import Symbol, UnresolvedSymbol



class  Section(object):
    def __init__(self, elf, order, name, _type, addr, offset, size):
        self.elf = elf
        self.order = order
        self.name = name
        self._type = _type
        self.addr = addr 
        self.offset = offset
        self.size = size        
        self._content = None
        
    @property
    def content(self):
        if self._content == None:
            with open(self.elf.objfile, "rb") as _f:
                _f.seek( self.offset )
                self._content = bytearray( _f.read( self.size ) )
        return self._content        


class ReadElfSource(object):
    def __init__(self, arch, objfile):
        self.objfile = objfile
        self.sections = None
        self.symbols = None
        self.arch = arch
        self._process()

    def _process(self):
        readelf_proc = subprocess.Popen(["/usr/bin/readelf","-a", self.objfile], stdout=subprocess.PIPE)
        stddata, stderr = readelf_proc.communicate()

        #
        # parse sections
        #        
        self.sections = dict( [(section.name, section) for section in self._parse_sections(stddata)] )

        self.symbol_tables = dict( [(sym_table, symbols) for sym_table, symbols in self._parse_symbol_tables(stddata)] )    
        #
        # use known symbol tables
        #
        syms = self.symbol_tables.get(".symtab", None)
        if syms == None:            
            syms = self.symbol_tables.get(".dynsym", None)
            if syms == None:
                raise Exception("No symbols in %s." % self.objfile)                        
        self.symbols = dict( [(sym.name, sym) for sym in syms] )        
        #
        # parse relocations
        #        
        self.relocations = []        
        for relocation in self._parse_relocations(stddata):
            # resolve symbols 
            relocation.symbol = self.symbols.get(relocation.symbol,  UnresolvedSymbol(relocation.symbol) )    
            # add relocations
            self.relocations.append( relocation )

    def _parse_sections(self, data):
        sections_expr = re.compile(r"Section Headers:(?:.*?)Al\n(.*?)Key to Flags:", re.DOTALL|re.M)
        m = sections_expr.search( data )
        if m:
            section_expr = re.compile(r"\[\s*(?P<order>[0-9a-fA-F]+)\]\s(?P<name>.*?)\s(?:\s*)(?P<type>\w+)\s+(?P<addr>[0-9a-fA-F]+)\s(?P<offset>[0-9a-fA-F]+)\s(?P<size>[0-9a-fA-F]+)")
            for line in m.groups()[0].split("\n"):
                m = section_expr.search(line)
                if m:
                    order, name, _type, addr, offset, size = m.groups()
                    yield Section(self, int(order), name, _type, int(addr,16), int(offset,16), int(size,16))
                
        
    def _parse_relocations(self, data):
        while len(data) > 0:
            relocation_begin_expr = re.compile(r"Relocation section '(.*?)' at offset (0x[0-9a-fA-F]+) contains (\d+) entries:(?:.*?)(?:Sym\. Name.*?\n)(.*?)\n\n", re.DOTALL|re.M)
            m = relocation_begin_expr.search( data )
            if m:
                rel_entries_count  = int( m.groups()[2] )
                rel_offset = int( m.groups()[1], 16)
                                
                rel_section_name = m.groups()[0]                
                
                if not rel_section_name.startswith(".rel"):
                    raise Exception("'%s' non .rel* relocation section.(?)" % rel_section_name)
                    
                rel_section_name = rel_section_name[5:]
                rel_entries_data = m.groups()[3]
                rel_entry_expr = re.compile(r"(?P<offset>[0-9a-fA-F]+)\s+(?P<info>[0-9a-fA-F]+)\s+(?P<type>.+?)\s(?P<symvalue>.*)")
                
                for offset, info, reltype, values in rel_entry_expr.findall(rel_entries_data):
                    symvalues = values.split()  
                    yield Relocation.byType(reltype)(self, self.arch, reltype, int(offset,16), int(info,16), symvalues )

                data = data[m.end():]
            else:
                break

    def _parse_symbol_tables(self, data):
        while len(data) > 0:
            sym_table_expr = re.compile(r"Symbol table '(.*?)' (?:.*?)Name\n\s+(.*?)\n\n", re.DOTALL|re.M)
            m = sym_table_expr.search(data)
            if m:
                symbol_table_name = m.groups()[0]
                symbol_table_data = m.groups()[1]

                sym_entry_expr = re.compile(r"(?P<num>\d+):\s+(?P<value>[0-9a-fA-F]+)\s+(?P<size>\d+)\s(?P<type>\w+)\s+(?P<bind>\w+)\s+(?P<visibility>\w+)\s+(?P<ndx>.*?)\s(?P<name>.*)")
                symbols = []
                for num, value, size, _type, binding, visibility, ndx, name in sym_entry_expr.findall( symbol_table_data ):
                    symbols.append( Symbol(int(num), int(value,16), int(size), _type, binding, visibility, ndx, name) )
                yield symbol_table_name, symbols
                data = data[m.end():]
            else:
                break
